#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Create by zhang
# Create on 2022/8/9 17:26
from typing import List, Tuple

import pandas as pd
from pandas import DataFrame
from numpy import nan, array, ndarray

from core.index_base import IndexBase
from domain.transaction_data.common.column_name import StockTradeDataColumnName


class Holt(IndexBase):
    """
    霍尔特的线性趋势法
    Holt (1957) 将简单指数平滑法扩展到可以预测具有趋势的数据。该方法包含一个预测方程和两个平滑方程（一个用于水平，一个用于趋势）
    Forecast equation: \[{y_{t + h|t}} = {\ell _t} + h{b_t}\]
    Level equation: \[{\ell _t} = \alpha {y_t} + (1 - \alpha )({\ell _{t - 1}} + {b_{t - 1}})\]
    Trend equation: \[{b_t} = {\beta ^*}({\ell _t} - {\ell _{t - 1}}) + (1 - {\beta ^*}){b_{t - 1}}\]
    其中\ell _t表示t时刻该时间序列的水平的估计值，b_t表示该时间序列在t时刻的趋势（斜率）的估计，\alpha是水平\[0 \le \alpha  \le 1\]的平滑参数，
    \beta ^*是趋势\[0 \le {\beta ^ * } \le 1\]的平滑参数。
    与简单的指数平滑一样，这里的水平方程表明\ell _t是观测值y_t和t时刻的向前一步训练预测值的加权平均值，这里由\ell _{t - 1}给出。
    趋势方程表明，b_t是基于{\ell _t} - {\ell _{t - 1}}和前一个趋势的估计值b_{t - 1}在t时刻的估计值的加权平均值。
    """

    def __init__(self, period=5, key: str = StockTradeDataColumnName.CLOSE):
        super(IndexBase, self).__init__()
        self.period = period
        self.key = key

    def compute(self, data: DataFrame):
        sma_key = f"{self.key}_holt_{self.period}"
        self.added_keys.append(sma_key)
        data[sma_key] = nan
        data[sma_key] = data[self.key].rolling(self.period).mean()